{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "8bda5ed3",
      "metadata": {
        "id": "8bda5ed3"
      },
      "source": [
        "# Setup\n",
        "\n",
        "## Install dependencies:\n",
        "\n",
        "```\n",
        "pip install ftfy regex tqdm autofaiss tensorflow google-colab\n",
        "pip install git+https://github.com/openai/CLIP.git\n",
        "```\n",
        "\n",
        "Dependending on if you want to GPU-accelerate the nearest-neighbors index:\n",
        "\n",
        "```pip install faiss-gpu``` OR ```pip install faiss-cpu```\n",
        "\n",
        "## Run setup script\n",
        "\n",
        "Run `setup.sh`.\n",
        "\n",
        "It will create the following file structure:\n",
        "\n",
        "```\n",
        "laion400m/\n",
        "  |- train/\n",
        "      |- metadata\n",
        "          |- metadata_0.parquet\n",
        "          |- ...\n",
        "          |- metadata_99.parquet\n",
        "      |- npy\n",
        "          |- img_emb_0.npy\n",
        "          |- ...\n",
        "          |- img_emb_99.npy\n",
        "      |- image.index\n",
        "  |- eval/\n",
        "      |- metadata\n",
        "          |- metadata_100.parquet\n",
        "          |- ...\n",
        "          |- metadata_199.parquet\n",
        "      |- npy\n",
        "          |- img_emb_100.npy\n",
        "          |- ...\n",
        "          |- img_emb_199.npy\n",
        "  |- al/\n",
        "      |- metadata\n",
        "          |- metadata_200.parquet\n",
        "          |- ...\n",
        "          |- metadata_299.parquet\n",
        "      |- npy\n",
        "          |- img_emb_200.npy\n",
        "          |- ...\n",
        "          |- img_emb_299.npy\n",
        "```\n",
        "\n",
        "Licensed under the Apache License, Version 2.0.",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ibSYp7Ced1-D",
      "metadata": {
        "cellView": "form",
        "id": "ibSYp7Ced1-D"
      },
      "outputs": [],
      "source": [
        "#@title Dependencies\n",
        "\n",
        "%load_ext autoreload\n",
        "%autoreload 2\n",
        "\n",
        "from annotator import Annotator\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "from clip_knn_service import create\n",
        "import utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "37f85c5e",
      "metadata": {
        "cellView": "form",
        "id": "37f85c5e"
      },
      "outputs": [],
      "source": [
        "#@title FAISS setup.\n",
        "knn_service = create(\n",
        "    indices_paths=\"indices_train.json\",\n",
        "\n",
        "    # Enable hdf5 caching for the metadata. This reduces metadata memory\n",
        "    # use to near zero, but on the first run will take a bit of time to\n",
        "    # create the memory-mapped hdf5 files.\n",
        "    enable_hdf5=True,\n",
        "\n",
        "    # Use an index with memory mapping, decreasing memory use to zero.\n",
        "    enable_faiss_memory_mapping=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b3ff95d0-463d-4765-a402-7c512b87e1b2",
      "metadata": {
        "cellView": "form",
        "id": "b3ff95d0-463d-4765-a402-7c512b87e1b2"
      },
      "outputs": [],
      "source": [
        "#@markdown ## Step 1: User Input\n",
        "#@markdown What visual concept would you like to classify? Describe it in a short phrase.\n",
        "concept = 'gourmet tuna' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Great! Next, please enter a few positive and negative queries for your concept\n",
        "#@markdown\n",
        "#@markdown Positive queries are short phrases that are synonyms or define a specific subset of the concept.\n",
        "#@markdown Negative queries are short phrases that define categories that are out-of-scope for your concept,\n",
        "#@markdown but are related. For example, if your concept is `gourmet tuna`, a positive\n",
        "#@markdown query could be `seared tuna` and a negative query could be\n",
        "#@markdown `canned tuna`. You do not have to fill out all the queries, but try to think of at least one positive and negative (the more the better).\n",
        "#@markdown Our system uses commas to separate different queries. For example if you want to use `tuna sushi` and `seared tuna` as postive queries, enter them as `tuna sushi, seared tuna`.\n",
        "#@markdown You do not need to re-enter the concept as a query.\n",
        "\n",
        "#@markdown Positive queries:\n",
        "positive_queries = 'tuna sushi, seared tuna, tuna sashimi'  #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Negative queries:\n",
        "negative_queries = 'canned tuna, tuna sandwich, tuna fish'  #@param {type:\"string\"}\n",
        "\n",
        "def convert_query_string_to_list(query_str):\n",
        "  query_list = query_str.split(',')\n",
        "  # Remove starting spaces.\n",
        "  query_list = [x.lstrip(' ') for x in query_list]\n",
        "  # Remove empty strings\n",
        "  return [x for x in query_list if x]\n",
        "\n",
        "queries = convert_query_string_to_list(positive_queries)\n",
        "queries.extend(convert_query_string_to_list(negative_queries))\n",
        "queries.append(concept)\n",
        "\n",
        "round_num = 0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "-WTMzSshe5ml",
      "metadata": {
        "cellView": "form",
        "id": "-WTMzSshe5ml"
      },
      "outputs": [],
      "source": [
        "#@markdown ## Step 2: Initial User Labeling\n",
        "#@markdown Using your concept and provided queries, we find relevant images and ask you to rate them. We bring up a rating UI. Here are the instructions:\n",
        "#@markdown 1. You can click the \"Positive\", \"Negative\", or \"No Image\" buttons to rate the image. We do not train on \"No Image\" ratings.\n",
        "#@markdown 2. Alternatively you can press the '1', '2', '3' keys as keyboard shortcuts for positive, negative, and no image respectively. However, to enable keyboard shortcuts, make sure that you have the rating UI focused -- an easy way to ensure this is to simply click on the image with your mouse.\n",
        "#@markdown 3. In case you want to re-rate images, you can use the left and right arrow keys to move forward and backward images. Simply press the buttons or use the '1', '2', '3' keys to rate again.\n",
        "\n",
        "#@markdown Finding images and displaying the rating interface takes a few seconds.\n",
        "\n",
        "# Randomly (but deterministically) take 50 images for initial labeling.\n",
        "SAMPLE_SIZE = 100\n",
        "nb_list = []\n",
        "\n",
        "for query in queries:\n",
        "  nb_list += knn_service.query(\n",
        "        text_input=query,\n",
        "        modality=\"image\",\n",
        "        num_images=100,\n",
        "        num_result_ids=100,\n",
        "        deduplicate=True,\n",
        "    )\n",
        "np.random.seed(0)\n",
        "np.random.shuffle(nb_list)\n",
        "\n",
        "# Take double the sample size initially because some images won't exist.\n",
        "nb_list = nb_list[:SAMPLE_SIZE*2]\n",
        "urls = [n['url'] for n in nb_list]\n",
        "\n",
        "# Make sure that the urls actually exist.\n",
        "exists = utils.check_images_exist_parallel(urls)\n",
        "\n",
        "# Filter by results where URL exist.\n",
        "exists_idx = np.argwhere(exists).squeeze()[:SAMPLE_SIZE]\n",
        "\n",
        "nb_list = [nb_list[i] for i in exists_idx]\n",
        "urls = [urls[i] for i in exists_idx]\n",
        "\n",
        "cache_path = None # Set this if you want to save ratings.\n",
        "annotator = Annotator(urls, cache_file_name=cache_path)\n",
        "annotations = annotator.annotate()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown ## Step 3: Train an Initial Model\n",
        "#@markdown Using the data you've rated, we train an initial model. This part should take about 2-3 minutes to run. You should see the training progress printing out.\n",
        "\n",
        "# Filter our the 'no image' ratings.\n",
        "annotations, nb_list = zip(*((anno, nb) for anno, nb in zip(annotations, nb_list) if anno != 'no image'))\n",
        "\n",
        "x_labeled = np.asarray([n['embedding'] for n in nb_list])\n",
        "y_labeled = np.asarray([anno == 'positive' for anno in annotations])\n",
        "\n",
        "# Collect data\n",
        "x_rand = np.load('laion400m/al/npy/img_emb_200.npy')\n",
        "x_rand = x_rand[:500000]\n",
        "y_rand = np.zeros((len(x_rand),))\n",
        "\n",
        "print('Got', len(x_labeled), 'ground-truth labels and', len(x_rand), 'random labels')\n",
        "\n",
        "# Upsample known positive and negatives.\n",
        "pos_features = x_labeled[y_labeled]\n",
        "pos_labels = y_labeled[y_labeled]\n",
        "\n",
        "ids = np.arange(len(pos_features))\n",
        "choices = np.random.choice(ids, len(x_rand))\n",
        "\n",
        "res_pos_features = pos_features[choices]\n",
        "res_pos_labels = pos_labels[choices]\n",
        "\n",
        "neg_features = x_labeled[~y_labeled]\n",
        "neg_labels = y_labeled[~y_labeled]\n",
        "\n",
        "ids = np.arange(len(neg_features))\n",
        "choices = np.random.choice(ids, len(x_rand)//2)\n",
        "res_neg_features = neg_features[choices]\n",
        "res_neg_labels = neg_labels[choices]\n",
        "\n",
        "# Combine labeled and random annotations.\n",
        "x_all = np.concatenate((res_pos_features, res_neg_features, x_rand), axis=0)\n",
        "y_all = np.concatenate((res_pos_labels, res_neg_labels, y_rand), axis=0)\n",
        "\n",
        "# Train a one-layer model (no hidden layers) since we have so few labels.\n",
        "model = utils.create_classifier(layer_dims=[16])\n",
        "print('The model has', model.count_params(), 'parameters')\n",
        "\n",
        "utils.train_model(model, x_all, y_all, validation_data=0.0, verbose=1)\n",
        "round_num = 0"
      ],
      "metadata": {
        "cellView": "form",
        "id": "YG3yl8JT2RoC"
      },
      "id": "YG3yl8JT2RoC",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}